package com.yupi.project.utils;

import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 基于Redis的令牌桶算法实现流量控制工具类
 *
 * @author yupi
 */
@Component
@Slf4j
public class RateLimiterUtil {

    @Resource
    private StringRedisTemplate stringRedisTemplate;
    
    @Resource
    private DefaultRedisScript<Long> rateLimiterScript;
    
    @Resource
    private DefaultRedisScript<List> rateLimiterCasScript;

    private static final String RATE_LIMITER_KEY_PREFIX = "rate_limiter:";

    /**
     * 尝试获取令牌，实现流量控制
     *
     * @param key 限流的唯一标识，如接口ID或用户ID
     * @param maxTokens 令牌桶容量
     * @param tokensPerSecond 令牌生成速率（每秒）
     * @param timeout 获取令牌的超时时间（毫秒）
     * @return 是否获取到令牌
     */
    public boolean tryAcquire(String key, int maxTokens, double tokensPerSecond, long timeout) {
        String redisKey = RATE_LIMITER_KEY_PREFIX + key;
        long now = System.currentTimeMillis();
        long deadline = now + timeout;

        // 使用Redis的事务保证操作的原子性
        while (now <= deadline) {
            // 开启Redis事务
            stringRedisTemplate.multi();

            try {
                // 获取当前令牌数和上次更新时间
                String tokenCountStr = stringRedisTemplate.opsForValue().get(redisKey + ":tokens");
                String lastRefillTimeStr = stringRedisTemString userInterfaceKey = "user:" + userId + ":interface:" + interfaceId;plate.opsForValue().get(redisKey + ":lastRefillTime");

                double tokenCount = tokenCountStr != null ? Double.parseDouble(tokenCountStr) : maxTokens;
                long lastRefillTime = lastRefillTimeStr != null ? Long.parseLong(lastRefillTimeStr) : now;

                // 计算从上次更新到现在应该新增的令牌数
                double newTokens = (now - lastRefillTime) / 1000.0 * tokensPerSecond;
                tokenCount = Math.min(maxTokens, tokenCount + newTokens);

                // 如果令牌数大于等于1，则获取成功
                if (tokenCount >= 1) {
                    // 消耗一个令牌
                    tokenCount -= 1;
                    
                    // 更新令牌数和上次更新时间
                    stringRedisTemplate.opsForValue().set(redisKey + ":tokens", String.valueOf(tokenCount));
                    stringRedisTemplate.opsForValue().set(redisKey + ":lastRefillTime", String.valueOf(now));
                    
                    // 设置过期时间，避免长期占用Redis内存
                    long expireTime = (long) Math.ceil(maxTokens / tokensPerSecond) * 2;
                    stringRedisTemplate.expire(redisKey + ":tokens", expireTime, TimeUnit.SECONDS);
                    stringRedisTemplate.expire(redisKey + ":lastRefillTime", expireTime, TimeUnit.SECONDS);
                    
                    // 执行事务
                    stringRedisTemplate.exec();
                    return true;
                } else {
                    // 令牌不足，取消事务
                    stringRedisTemplate.discard();
                    
                    // 等待一段时间再尝试
                    try {
                        Thread.sleep(100);
                    } catch (InterruptedException e) {
                        Thread.currentThread().interrupt();
                        return false;
                    }
                }
            } catch (Exception e) {
                // 发生异常，取消事务
                stringRedisTemplate.discard();
                log.error("Rate limiter error", e);
                return false;
            }
            
            now = System.currentTimeMillis();
        }
        
        return false;
    }

    /**
     * 尝试获取令牌，使用默认超时时间
     *
     * @param key 限流的唯一标识
     * @param maxTokens 令牌桶容量
     * @param tokensPerSecond 令牌生成速率（每秒）
     * @return 是否获取到令牌
     */
    public boolean tryAcquire(String key, int maxTokens, double tokensPerSecond) {
        return tryAcquire(key, maxTokens, tokensPerSecond, 1000);
    }

    /**
     * 使用Lua脚本实现的令牌桶算法，保证原子性操作
     *
     * @param key 限流的唯一标识，如接口ID或用户ID
     * @param maxTokens 令牌桶容量
     * @param tokensPerSecond 令牌生成速率（每秒）
     * @param requestTokens 请求的令牌数量，默认为1
     * @return 是否获取到令牌
     */
    public boolean tryAcquireWithLua(String key, int maxTokens, double tokensPerSecond, int requestTokens) {
        try {
            String redisKey = RATE_LIMITER_KEY_PREFIX + key;
            long now = System.currentTimeMillis();
            
            // 执行Lua脚本
            List<String> keys = Collections.singletonList(redisKey);
            Object[] args = {String.valueOf(maxTokens), String.valueOf(tokensPerSecond), 
                           String.valueOf(now), String.valueOf(requestTokens)};
            
            Long result = stringRedisTemplate.execute(rateLimiterScript, keys, args);
            
            // 返回结果：1表示获取令牌成功，0表示获取令牌失败
            return result != null && result == 1L;
        } catch (Exception e) {
            log.error("Rate limiter with Lua script error", e);
            // 发生异常时，默认放行，避免影响正常业务
            return true;
        }
    }
    
    /**
     * 使用Lua脚本实现的令牌桶算法，请求单个令牌
     *
     * @param key 限流的唯一标识，如接口ID或用户ID
     * @param maxTokens 令牌桶容量
     * @param tokensPerSecond 令牌生成速率（每秒）
     * @return 是否获取到令牌
     */
    public boolean tryAcquireWithLua(String key, int maxTokens, double tokensPerSecond) {
        return tryAcquireWithLua(key, maxTokens, tokensPerSecond, 1);
    }
    
    /**
     * 使用基于CAS的Lua脚本实现的令牌桶算法，支持等待策略
     *
     * @param key 限流的唯一标识，如接口ID或用户ID
     * @param maxTokens 令牌桶容量
     * @param tokensPerSecond 令牌生成速率（每秒）
     * @param requestTokens 请求的令牌数量
     * @param maxWaitTime 最大等待时间（毫秒）
     * @return 是否获取到令牌
     */
    public boolean tryAcquireWithLuaAndWait(String key, int maxTokens, double tokensPerSecond, int requestTokens, long maxWaitTime) {
        try {
            String redisKey = RATE_LIMITER_KEY_PREFIX + key;
            long startTime = System.currentTimeMillis();
            long deadline = startTime + maxWaitTime;
            long now = startTime;
            
            // 版本号，用于CAS操作
            Integer version = null;
            
            // 退避策略参数
            int retryCount = 0;
            long baseWaitTime = 10; // 初始等待时间（毫秒）
            long maxBackoffTime = 1000; // 最大退避时间（毫秒）
            
            while (now <= deadline) {
                // 执行Lua脚本
                List<String> keys = Collections.singletonList(redisKey);
                Object[] args;
                
                if (version != null) {
                    // 使用版本号进行CAS操作
                    args = new Object[]{String.valueOf(maxTokens), String.valueOf(tokensPerSecond), 
                                String.valueOf(now), String.valueOf(requestTokens), String.valueOf(version)};
                } else {
                    // 首次尝试，不使用版本号
                    args = new Object[]{String.valueOf(maxTokens), String.valueOf(tokensPerSecond), 
                                String.valueOf(now), String.valueOf(requestTokens)};
                }
                
                List<Long> result = stringRedisTemplate.execute(rateLimiterCasScript, keys, args);
                
                if (result == null || result.size() < 4) {
                    log.error("Invalid result from rate limiter CAS script: {}", result);
                    return false;
                }
                
                // 解析结果
                Long status = result.get(0);          // 状态码：1成功，0失败
                Long currentTokens = result.get(1);   // 当前令牌数
                Long nextTokenTime = result.get(2);   // 下一个令牌预计可用时间（毫秒）
                Long currentVersion = result.get(3);  // 当前版本号
                
                // 更新版本号
                version = currentVersion.intValue();
                
                // 如果获取令牌成功，直接返回
                if (status == 1L) {
                    return true;
                }
                
                // 计算剩余等待时间
                long remainingTime = deadline - now;
                
                // 如果下一个令牌可用时间超过了剩余等待时间，直接返回失败
                if (nextTokenTime > remainingTime) {
                    log.warn("Next token time exceeds remaining time. Next token time: {}, Remaining time: {}", nextTokenTime, remainingTime);
                    return false;
                }
                
                // 使用指数退避策略计算等待时间
                long waitTime = Math.min(
                    Math.min(nextTokenTime, baseWaitTime * (1L << Math.min(retryCount, 10))), // 指数增长，但不超过2^10
                    maxBackoffTime // 不超过最大退避时间
                );
                
                // 确保等待时间不超过剩余时间
                waitTime = Math.min(waitTime, remainingTime);
                
                // 等待一段时间再重试
                if (waitTime > 0) {
                    try {
                        Thread.sleep(waitTime);
                    } catch (InterruptedException e) {
                        Thread.currentThread().interrupt();
                        return false;
                    }
                }
                
                // 增加重试计数
                retryCount++;
                now = System.currentTimeMillis();
            }
            
            // 超过最大等待时间，获取令牌失败
            return false;
        } catch (Exception e) {
            log.error("Rate limiter with CAS error", e);
            // 发生异常时，默认放行，避免影响正常业务
            return true;
        }
    }
    
    /**
     * 使用基于CAS的Lua脚本实现的令牌桶算法，请求单个令牌，支持等待策略
     *
     * @param key 限流的唯一标识，如接口ID或用户ID
     * @param maxTokens 令牌桶容量
     * @param tokensPerSecond 令牌生成速率（每秒）
     * @param maxWaitTime 最大等待时间（毫秒）
     * @return 是否获取到令牌
     */
    public boolean tryAcquireWithLuaAndWait(String key, int maxTokens, double tokensPerSecond, long maxWaitTime) {
        return tryAcquireWithLuaAndWait(key, maxTokens, tokensPerSecond, 1, maxWaitTime);
    }
    
    /**
     * 针对接口调用的流量控制，使用Lua脚本实现的令牌桶算法
     *
     * @param interfaceId 接口ID
     * @param userId 用户ID
     * @param maxTokens 令牌桶容量
     * @param tokensPerSecond 令牌生成速率（每秒）
     * @return 是否允许调用
     */
    public boolean isAllowedInterfaceRequest(Long interfaceId, Long userId, int maxTokens, double tokensPerSecond) {
        // 针对接口的全局流控
        String interfaceKey = "interface:" + interfaceId;
        boolean interfaceAllowed = tryAcquireWithLua(interfaceKey, maxTokens, tokensPerSecond);
        if (!interfaceAllowed) {
            log.warn("Interface rate limit exceeded for interface: {}", interfaceId);
            return false;
        }
        
        // 针对用户对特定接口的流控
        // 使用更合理的用户级别限流参数，确保单个用户不会占用过多资源，同时也不会过度限制
        // 假设系统预期同时服务的活跃用户数量为10个
        int expectedActiveUsers = 10;
        int userMaxTokens = Math.max(maxTokens / expectedActiveUsers, 5); // 确保每个用户至少有5个令牌
        double userTokensPerSecond = Math.max(tokensPerSecond / expectedActiveUsers, 1.0); // 确保每个用户至少有1个令牌/秒的速率
        
        String userInterfaceKey = "user:" + userId + ":interface:" + interfaceId;
        boolean userInterfaceAllowed = tryAcquireWithLua(userInterfaceKey, userMaxTokens, userTokensPerSecond);
        if (!userInterfaceAllowed) {
            log.warn("User rate limit exceeded for user: {} on interface: {}", userId, interfaceId);
            return false;
        }
        
        return true;
    }
    
    /**
     * 针对接口调用的流量控制，使用基于CAS的令牌桶算法，支持等待策略
     *
     * @param interfaceId 接口ID
     * @param userId 用户ID
     * @param maxTokens 令牌桶容量
     * @param tokensPerSecond 令牌生成速率（每秒）
     * @param maxWaitTime 最大等待时间（毫秒）
     * @return 是否允许调用
     */
    public boolean isAllowedInterfaceRequestWithWait(Long interfaceId, Long userId, int maxTokens, double tokensPerSecond, long maxWaitTime) {
        // 针对接口的全局流控
        String interfaceKey = "interface:" + interfaceId;
        boolean interfaceAllowed = tryAcquireWithLuaAndWait(interfaceKey, maxTokens, tokensPerSecond, maxWaitTime);
        if (!interfaceAllowed) {
            log.warn("Interface rate limit exceeded for interface: {} after waiting {} ms", interfaceId, maxWaitTime);
            return false;
        }
        
        // 针对用户对特定接口的流控
        // 使用更合理的用户级别限流参数，确保单个用户不会占用过多资源，同时也不会过度限制
        int expectedActiveUsers = 10;
        int userMaxTokens = Math.max(maxTokens / expectedActiveUsers, 5); // 确保每个用户至少有5个令牌
        double userTokensPerSecond = Math.max(tokensPerSecond / expectedActiveUsers, 1.0); // 确保每个用户至少有1个令牌/秒的速率
        
        // 计算用户级别的最大等待时间，通常应该小于全局等待时间
        long userMaxWaitTime = Math.min(maxWaitTime, 500); // 用户级别最多等待500ms
        
        String userInterfaceKey = "user:" + userId + ":interface:" + interfaceId;
        boolean userInterfaceAllowed = tryAcquireWithLuaAndWait(userInterfaceKey, userMaxTokens, userTokensPerSecond, userMaxWaitTime);
        if (!userInterfaceAllowed) {
            log.warn("User rate limit exceeded for user: {} on interface: {} after waiting {} ms", userId, interfaceId, userMaxWaitTime);
            return false;
        }
        
        return true;
    }
}